import torch
import torch.nn as nn
import torch.nn.functional as F

def bilinear_pooling(x, y):
    batch_size = x.size(0)
    channel_size = x.size(1)
    feature_size = x.size(2) * x.size(3)


    x = x.view(batch_size, channel_size, feature_size)
    y = y.view(batch_size, channel_size, feature_size)
    out = (torch.bmm(x, torch.transpose(y, 1, 2)) / feature_size).view(batch_size, -1)
    out = torch.nn.functional.normalize(torch.sign(out) * torch.sqrt(torch.abs(out) + 1e-10))

    return out  # [N,C*C]

if __name__ == "__main__":

    x1 = torch.randn(16, 768, 7, 7)
    x2 = torch.randn(16, 768, 7, 7)
    print("x1: ", x1.shape)

    x_b_p = bilinear_pooling(x1, x2)
    print("x_b_p: ", x_b_p.shape)


    # # 实现方式：pytorch
    # x = torch.randn(8, 512, 12, 12)
    # batch_size = x.size(0)
    # feature_size = x.size(2) * x.size(3)
    # x = x.view(batch_size, 512, feature_size)
    # x = (torch.bmm(x, torch.transpose(x, 1, 2)) / feature_size).view(batch_size, -1)
    # x = torch.nn.functional.normalize(torch.sign(x) * torch.sqrt(torch.abs(x) + 1e-10))
    # print("x: ", x.shape)



